#ifndef __LZJ_PARSER_CPP__
#define __LZJ_PARSER_CPP__

#include "parser_types.cpp"
#include <vector>
#include <iostream>
#include <string>
#include <unordered_map>
#include <set>
#include <algorithm>
#include <deque>

using std::deque;
using std::fstream;
using std::set;
using std::string;
using std::stringstream;
using std::unordered_map;
using std::vector;

static const int _startSymbol = 0;          // 分析栈内起始符号
static vector<string> _symbols;             // 语法符号
static int _symbolsTerminalOffset;          // _symbols[0, offset) 为非终结符，后面的为终结符
static int _endMarker;                      // _symbols数组长度
static unordered_map<int, set<int>> first;  // first集
static vector<RuleType> _rules;             // 规则
static vector<vector<LRItem>> _collection;  // 状态机状态，_collection[i]表示 i 状态下的项目集合
static vector<vector<Action>> _parseTable;  // 动作
static bool _parseTableHasConflict;         // LR表有冲突
static unordered_map<string, int> tokenName2Index;  // 将token名转换成index

// public APIs

void createParser(std::istream &);
PtrParseTreeNode lzjParser(std::istream &, deque<Token> (*)(std::istream &, std::ostream &));
void printParseTable(std::ostream &);
void printParseTree(std::ostream &, const PtrParseTreeNode &, const string &);

// private functions

static void clear();
static bool isNonTerminal(int);
static bool isTerminal(int);
static bool isTerminalOrEndMarker(int);
static bool setsEqual(vector<LRItem>, vector<LRItem>);
static void getFirst(int);
static vector<LRItem> closure(vector<LRItem>);
static vector<LRItem> goTo(vector<LRItem>, int);
static void printState(int);
static void createCollection();
static void addActionToParseTable(int, int, Action);
static int token2Symbol(const Token &);


static void clear() {
    _symbols.clear();
    _rules.clear();
    _collection.clear();
    _parseTable.clear();
    _parseTableHasConflict = false;
    // _parseSteps.clear();
}

static inline bool isNonTerminal(int symbol) {
    return symbol < _symbolsTerminalOffset && symbol> 0;
}

static inline bool isTerminal(int symbol) { 
    return symbol >= _symbolsTerminalOffset && symbol < _symbols.size() - 1;
}

static inline bool isTerminalOrEndMarker(int symbol) {
    return symbol >= _symbolsTerminalOffset && symbol < _symbols.size();
}

static bool setsEqual(vector<LRItem> s1, vector<LRItem> s2) {
    if (s1.size() != s2.size()) {
        return false;
    }
    sort(s1.begin(), s1.end());
    sort(s2.begin(), s2.end());
    for (int i = 0; i < s1.size(); ++i){
        if(s1[i] != s2[i])
            return false;
    }
    return true;
}

static void getFirst(int symbol) {
    if(first.count(symbol)){
        return;
    }
    if (isTerminalOrEndMarker(symbol)) {
        first.insert({symbol, {symbol}});
        return;
    }
    set<int> f;
    for (auto i: _rules) {
        if (i.lhs == symbol && !i.rhs.empty() && i.rhs[0] != symbol) {
            getFirst(i.rhs[0]);
            f.insert((first[i.rhs[0]]).begin(), first[i.rhs[0]].end());
        }
    }
    first.insert({symbol, f});
};

static vector<LRItem> closure(vector<LRItem> s) {
    // cout << "Closure: ";
    // for(auto i: s){
    //     cout << _symbols[i.rule.lhs) << ", ";
    // }
    // cout << endl;
    vector<LRItem> closureSet = s;
    bool newItemAdded = true;
    while (newItemAdded) {
        newItemAdded = false;
        for (int i = 0; i < closureSet.size(); ++i) {
            const int symbol = closureSet[i].getExpectedSymbol();
            for (auto rule: _rules) {
                if (rule.lhs == symbol) { // 找到拥有该产生式的规则
                    set<int> fst =
                        first[closureSet[i].position + 1 < closureSet[i].rule.rhs.size()
                                     ? closureSet[i].rule.rhs[closureSet[i].position + 1]
                                     : closureSet[i].lookahead];
                    for (auto k: fst) {
                        LRItem item(rule, 0, k);
                        bool newItem = true;
                        if(!contains(closureSet, item)){
                            closureSet.push_back(item);
                            newItemAdded = true;
                        }
                    }
                }
            }
        }
    }
    return closureSet; // TODO many set
};

static vector<LRItem> goTo(vector<LRItem> s, int symbol) {
    vector<LRItem> gotoSet;
    for (auto i: s) {
        if (i.position < i.rule.rhs.size()
            && i.rule.rhs[i.position] == symbol) {
            gotoSet.push_back(LRItem(i.rule, i.position + 1, i.lookahead));
            // cout << _symbols[gotoSet.back().rule.lhs) << endl;
        }
    }
    // cout << "----\n";
    return closure(gotoSet);
}

static void printState(int i) {
    cout << "State " << i << "\n";
    for (auto j : _collection[i]) {
        cout << _symbols[j.rule.lhs] << " -> ";
        for (int k = 0; k < j.rule.rhs.size(); ++k) {
            if (k == j.position) {
                cout << ".";
            }
            cout << _symbols[j.rule.rhs[k]] << " ";
        }
        if (j.position == j.rule.rhs.size())
            cout << ".";
        cout << " @ " << _symbols[j.lookahead];
        cout << endl;
    }
    cout << endl << endl;
}

static void createCollection() {
    cout << "Creating parser...\n";
    LRItem initialItem = LRItem(_rules[0], 0, _endMarker);
    _collection.push_back(closure({initialItem}));
    bool newSetAdded = true;
    while (newSetAdded) {
        newSetAdded = false;
        for (int i = 0; i < _collection.size(); ++i) {
            for (int symbol = 1; symbol < _symbols.size() - 1; ++symbol) {
                vector<LRItem> nextState = goTo(_collection[i], symbol);
                if (nextState.size() != 0) {
                    bool newSet = true;
                    for (int j = 0; j < _collection.size(); ++j) {
                        if (setsEqual(_collection[j], nextState)) {
                            newSet = false;
                            break;
                        }
                    }
                    if (newSet) {
                        _collection.push_back(nextState);
                        // cout <<"From "<< i<<  " symbol: " << _symbols[symbol) << endl;
                        // printState(_collection.size() - 1);
                        newSetAdded = true;
                    }
                }
            }
        }
    }
}

static void addActionToParseTable(int state, int symbol, Action action) {
    if (_parseTable[state][symbol].atype == ActionType::UNDEFINED) { // 无冲突
        _parseTable[state][symbol] = action;
    } else if (_parseTable[state][symbol].atype == ActionType::CONFLICT) { // 有冲突
        bool newAction = true;
        for (int i = 0; i < _parseTable[state][symbol].actions.size(); ++i) {
            if (*_parseTable[state][symbol].actions[i] == action) {
                newAction = false;
                break;
            }
        }
        if (newAction) {
            _parseTable[state][symbol].addAction(action);
        }
    } else if (_parseTable[state][symbol] != action) { // 有冲突
        _parseTable[state][symbol].addAction(action);
        _parseTableHasConflict = true;
    }
}

static void createParseTable() {
    _parseTable = vector<vector<Action>>(_collection.size(), vector<Action>(_symbols.size()));
    _parseTableHasConflict = false;
    for (int i = 0; i < _collection.size(); ++i) {
        for (int j = 0; j < _collection[i].size(); ++j) { // 遍历每一个项目查看它将要接受的符号
            const int symbol = _collection[i][j].getExpectedSymbol();
            // cout << i << ", " << j << endl;
            if (symbol == -1) { // 到尾了，应reduce或accept
                if (_collection[i][j].rule.lhs != _startSymbol) { // reduce
                    addActionToParseTable(i,
                                          _collection[i][j].lookahead,
                                          Action(_collection[i][j].rule));
                } else if (_collection[i][j].lookahead == _endMarker) { // accept
                    _parseTable[i][_endMarker] = Action(true);
                }
            }
            if (isTerminal(symbol)) { // shift
                vector<LRItem> gotoSet = goTo(_collection[i], symbol);
                for (int k = 0; k < _collection.size(); ++k) {
                    if (setsEqual(_collection[k], gotoSet)) {
                        addActionToParseTable(i, symbol, Action(k));
                    }
                }
            }
        }
        for (int symbol = 1; symbol < _symbolsTerminalOffset; ++symbol) {
            vector<LRItem> gotoSet = goTo(_collection[i], symbol);
            for (int k = 0; k < _collection.size(); ++k) {
                if (setsEqual(_collection[k], gotoSet)) {
                    _parseTable[i][symbol] = k;
                }
            }
        }
    }
}

void createParser (std::istream& ss) {
    clear();
    vector<vector<string>> rulesSymbols;
    vector<string> lines;
    string tmp;
    while(std::getline(ss, tmp)){
        tmp = trim(tmp);
        if(tmp.length()!=0){
            tmp = replace(tmp, "->", " ");
            lines.push_back(tmp);
        }
    }
    for (int i = 0; i < lines.size(); ++i) {
        rulesSymbols.push_back(split(trim(lines[i])));
        if(rulesSymbols.back().empty())
            rulesSymbols.pop_back();
        if (!contains(_symbols, rulesSymbols[i][0])) {
            _symbols.push_back(rulesSymbols[i][0]);
        }
    }
    _symbolsTerminalOffset = _symbols.size();
    for (auto i: rulesSymbols) {
        for (auto j: i) {
            if (!contains(_symbols, j) && j!="->") {
                _symbols.push_back(j);
            }
        }
    }
    _endMarker = _symbols.size();
    _symbols.push_back("$");
    for (int i = 0; i < rulesSymbols.size(); ++i) {
        _rules.push_back(RuleType(i));
        cout << i << ": ";
        for (int j = 0; j < rulesSymbols[i].size(); ++j) {
            cout << "\'" << rulesSymbols[i][j] << "\' ";
            if (j == 0) {
                _rules[i].lhs = indexOf(_symbols, rulesSymbols[i][j]);
            } else {
                _rules[i].rhs.push_back(indexOf(_symbols, rulesSymbols[i][j]));
            }
        }
        cout << endl;
    }
    RuleType::ruleNames = _symbols;//TODO stupid
    cout << "FF";
    for (int i = 0; i < _symbols.size(); ++i) {
        getFirst(i);
    }
    for(auto i:first){
        cout << "First of " << _symbols[i.first] << ": ";
        for(auto j:i.second){
            cout << _symbols[j] << " ";
        }
        cout << endl;
    }
    createCollection();
    createParseTable();
    cout << "Finish!\nParse table size: " << _parseTable.size() << "*" << _parseTable[0].size()
         << endl << "Parser created successfully!!!\n";
}

static int token2Symbol(const Token &token) {
    string tokenName = getTokenLiteral(token.tokenType);
    if (tokenName2Index.count(tokenName)) {
        return tokenName2Index[tokenName];
    }else{
        int index = indexOf(_symbols, tokenName);
        if(index == -1){
            cout << "[Error] Token \'" << tokenName
                 << "\' did not exist in grammar!"
                 << "Please check your grammar definition file!\n"
                 << "Parser terminated.\n";
            exit(0);
        }
        tokenName2Index[tokenName] = index;
        return index;
    }
}

PtrParseTreeNode lzjParser(std::istream& inputS,  deque<Token> input)
{
    if (_parseTableHasConflict) {
        // _parseSteps.push_back(ParseStepLog({},{},{},
        //     Action(ParseException(0,0,"Conflict in parse table"))));
        cout << "Conflict in parse table!\nParser terminated.\n";
        exit(0);
    }
    stack<int> stateStack;
    stack<PtrParseTreeNode> symbolStack;
    Action action;
    stateStack.push(0);
    while (true) {
        action = _parseTable[stateStack.top()][token2Symbol(input[0])];
        if (action.atype == ActionType::UNDEFINED) {
            cout << "[Error] line " << input[0].startPos.line << "(" << input[0].startPos.col
                 << "): \'" << input[0].word << "\'\nSyntax error!";
            exit(0);
        }
        if(action.atype == ActionType::SHIFT){
            symbolStack.push(make_shared<ParseTreeNode>(input.front()));
            input.pop_front();
            stateStack.push(action.state);
        } else if(action.atype == ActionType::REDUCE) {
            deque<PtrParseTreeNode> children; // 如果规约，则把所有归约的字符加进这里
            for (int i = 0; i < action.rule.rhs.size(); ++i) {
                children.push_front(symbolStack.top());
                symbolStack.pop();
                stateStack.pop();
            }
            PtrParseTreeNode node = make_shared<ParseTreeNode>(action.rule, children);
            symbolStack.push(node);
            // GOTO
            stateStack.push(_parseTable[stateStack.top()][action.rule.lhs].state);
        } else if(action.atype == ActionType::ACCEPT) {
            break;
        } else {
            cout << "Syntax error" << endl;
            exit(0);
        }
    }
    return symbolStack.top();
}

void printParseTree(std::ostream &oss, const PtrParseTreeNode& t, const string& indent = "") {
    oss << indent << "|-- ";
    if (t->isTerminal()) {
        oss << t->token.word << endl;
    } else {
        oss << _symbols[t->ruleT.lhs] << endl;
        for (auto i : t->children) {
            printParseTree(oss, i, indent + "  ");
        }
    }
}

void printParseTable(std::ostream& grammarOut){
    grammarOut << "state, ";
    for (int i = 1; i < _symbols.size(); ++i) {
        grammarOut << _symbols[i] << ", ";
    }
    grammarOut << endl;
    for (int i = 0; i < _parseTable.size(); ++i) {
        grammarOut << i << ",";
        for (int j = 1; j < _parseTable[i].size(); ++j) {
            grammarOut << _parseTable[i][j] << ", ";
        }
        grammarOut << endl;
    }
}

void printAST(std::ostream &oss, const PtrASTNode &t, const string &indent = "") {
    oss << indent << "|-- " << t->name << "{ ";
    for(auto i: t->fields){
        oss << i.first << ": ";
        if(i.second.empty())
            oss << i.second[0].word;
        else{
            for(auto j: i.second){
                oss << j.word << " ";
            }
        }
        oss << ", ";
    }
    oss << " }" << endl;
    for(auto i: t->children){
        printAST(oss, i, indent + "  ");
    }
}

#endif